"""
加密货币相关的Celery任务
"""
import logging
import datetime
from celery import group
from celery.utils.log import get_task_logger
from sqlalchemy import func

from modules.tasks.celery_app import app
from modules.models.base import get_db
from modules.models.stock import SecurityInfo, CollectionTask
from modules.data_collector import collector

# 配置任务日志
logger = get_task_logger(__name__)


@app.task(bind=True, queue='etf_data', max_retries=3, retry_backoff=60)
def update_crypto_list(self, limit=200):
    """更新加密货币列表信息"""
    logger.info(f"开始更新加密货币列表信息，数量限制: {limit}")
    try:
        new_count, update_count = collector.collect_market_data(
            "crypto_list", 
            limit=limit
        )
        logger.info(f"加密货币列表更新完成: 新增 {new_count} 种，更新 {update_count} 种")
        return {
            "status": "success",
            "new_count": new_count,
            "update_count": update_count
        }
    except Exception as e:
        logger.error(f"更新加密货币列表失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='etf_data', max_retries=3, retry_backoff=60)
def collect_crypto_data(self, symbol, start_date=None, end_date=None):
    """
    采集单个加密货币的历史数据
    
    Args:
        symbol: 加密货币代码
        start_date: 开始日期, 格式为YYYY-MM-DD
        end_date: 结束日期, 格式为YYYY-MM-DD
    """
    logger.info(f"开始采集加密货币 {symbol} 的历史数据")
    
    try:
        # 创建任务记录
        db = next(get_db())
        
        # 转换日期格式（如果需要）
        start_date_obj = None
        end_date_obj = None
        
        if start_date:
            try:
                start_date_obj = datetime.datetime.strptime(start_date, '%Y-%m-%d').date()
            except ValueError:
                start_date_obj = datetime.datetime.strptime(start_date, '%Y%m%d').date()
                start_date = start_date_obj.strftime('%Y-%m-%d')
        
        if end_date:
            try:
                end_date_obj = datetime.datetime.strptime(end_date, '%Y-%m-%d').date()
            except ValueError:
                end_date_obj = datetime.datetime.strptime(end_date, '%Y%m%d').date()
                end_date = end_date_obj.strftime('%Y-%m-%d')
        
        task_record = CollectionTask(
            task_id=self.request.id,
            task_type='crypto_daily',
            security_code=symbol,
            start_date=start_date_obj,
            end_date=end_date_obj,
            status='pending'
        )
        db.add(task_record)
        db.commit()
        db.close()
        
        # 执行数据采集
        count = collector.collect_market_data(
            "crypto_daily", 
            code=symbol, 
            start_date=start_date, 
            end_date=end_date, 
            task_id=self.request.id
        )
        
        logger.info(f"加密货币 {symbol} 采集完成，获取 {count} 条数据")
        
        # 如果有新数据，触发技术指标计算
        if count > 0:
            # 使用链式任务计算技术指标
            from modules.tasks.stock_tasks import calculate_stock_indicators
            calculate_stock_indicators.delay(symbol)
        
        return {
            "status": "success",
            "symbol": symbol,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集加密货币 {symbol} 数据失败: {str(e)}")
        # 更新任务记录状态
        try:
            db = next(get_db())
            task = db.query(CollectionTask).filter(CollectionTask.task_id == self.request.id).first()
            if task:
                task.status = 'failed'
                task.error_message = str(e)
                task.update_time = datetime.datetime.now()
                db.commit()
        except Exception as db_e:
            logger.error(f"更新任务状态失败: {str(db_e)}")
        finally:
            db.close()
            
        raise self.retry(exc=e)


@app.task(bind=True, queue='etf_data')
def schedule_daily_crypto_update(self, top_limit=50):
    """
    调度每日加密货币数据更新任务
    每天执行，获取市值排名靠前的加密货币的最新数据
    
    Args:
        top_limit: 更新市值排名前多少的加密货币，默认50
    """
    logger.info(f"开始调度每日加密货币数据更新任务，数量限制: {top_limit}")
    
    try:
        # 首先更新加密货币列表
        update_crypto_list.delay(limit=200)  # 采集更多，但只更新前top_limit个
        
        # 获取当前日期和前30天
        end_date = datetime.datetime.now().strftime('%Y-%m-%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y-%m-%d')
        
        # 获取活跃加密货币
        db = next(get_db())
        cryptos = db.query(SecurityInfo).filter(
            SecurityInfo.type == 'crypto',
            SecurityInfo.is_active == 1
        ).order_by(
            # 按ID排序，假设ID越小市值越大（实际可能需要按其他字段排序）
            SecurityInfo.id
        ).limit(top_limit).all()
        db.close()
        
        logger.info(f"准备更新 {len(cryptos)} 种加密货币的历史数据")
        
        # 创建并行任务组
        tasks = []
        for crypto in cryptos:
            tasks.append(collect_crypto_data.s(crypto.code, start_date, end_date))
        
        # 分批执行任务，避免一次创建过多任务
        batch_size = 10
        for i in range(0, len(tasks), batch_size):
            batch_tasks = tasks[i:i+batch_size]
            group(batch_tasks).apply_async()
            logger.info(f"已调度 {i+len(batch_tasks)}/{len(tasks)} 种加密货币数据采集任务")
        
        return {
            "status": "success",
            "message": f"已调度 {len(cryptos)} 种加密货币的数据更新任务"
        }
    except Exception as e:
        logger.error(f"调度加密货币更新任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='etf_data')
def collect_crypto_historical_data(self, symbol, days=365):
    """
    采集加密货币的历史数据
    
    Args:
        symbol: 加密货币代码
        days: 采集最近多少天的数据，默认365天
    """
    logger.info(f"开始采集加密货币 {symbol} 的 {days} 天历史数据")
    
    try:
        # 计算起始日期
        end_date = datetime.datetime.now().strftime('%Y-%m-%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=days)).strftime('%Y-%m-%d')
        
        # 执行数据采集
        result = collect_crypto_data.delay(symbol, start_date, end_date)
        
        return {
            "status": "success",
            "message": f"已调度加密货币 {symbol} 的历史数据采集任务",
            "task_id": result.id
        }
    except Exception as e:
        logger.error(f"调度加密货币历史数据采集任务失败: {str(e)}")
        raise


# 测试函数
def test_crypto_tasks():
    """测试加密货币任务函数"""
    # 更新加密货币列表
    update_result = update_crypto_list.delay(limit=20)
    print(f"更新加密货币列表任务已提交，任务ID: {update_result.id}")
    
    # 等待任务完成
    result = update_result.get(timeout=120)
    print(f"更新加密货币列表结果: {result}")
    
    # 采集单个加密货币数据
    db = next(get_db())
    crypto = db.query(SecurityInfo).filter(
        SecurityInfo.type == 'crypto'
    ).first()
    db.close()
    
    if crypto:
        # 计算最近30天的日期范围
        end_date = datetime.datetime.now().strftime('%Y-%m-%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y-%m-%d')
        
        collect_result = collect_crypto_data.delay(crypto.code, start_date, end_date)
        print(f"加密货币 {crypto.code} 数据采集任务已提交，任务ID: {collect_result.id}")
        
        # 等待任务完成
        result = collect_result.get(timeout=120)
        print(f"加密货币数据采集结果: {result}")
    else:
        print("没有找到加密货币数据，无法测试采集任务")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_crypto_tasks()